package container

type Set[K comparable] struct {
	keyToIdx map[K]int
	items    []K
}

func NewSet[K comparable]() *Set[K] {
	return &Set[K]{keyToIdx: map[K]int{}, items: make([]K, 0, 4)}
}

func (s *Set[K]) ForEach(f func(key K) bool) {
	for index := range s.items {
		if index >= s.Cnt() {
			break
		}
		if !f(s.items[index]) {
			break
		}
	}
}

func (s *Set[K]) Has(key K) bool {
	_, ok := s.keyToIdx[key]
	return ok
}

func (s *Set[K]) Cnt() int {
	return len(s.keyToIdx)
}

func (s *Set[K]) Add(key K) bool {
	_, ex := s.keyToIdx[key]
	if ex {
		return false
	}
	curLen := len(s.items)
	s.checkSize(curLen + 1)
	s.keyToIdx[key] = curLen
	s.items = append(s.items, key)
	return true
}

func (s *Set[K]) Remove(key K) bool {
	idx, ok := s.keyToIdx[key]
	if !ok {
		return false
	}
	delete(s.keyToIdx, key)
	curLen := len(s.items)
	tailValue := s.items[curLen-1]
	s.items[idx] = tailValue
	s.keyToIdx[tailValue] = idx
	s.items = s.items[0 : curLen-2]
	s.checkSize(len(s.items))
	return true
}

func (s *Set[K]) checkSize(minSize int) {
	caps := cap(s.items)
	if caps < minSize*2 && caps > minSize {
		return
	}
	minAllocateSize := s.minAllocSize(minSize)
	newcap := make([]K, minAllocateSize)
	copy(newcap, s.items)
	s.items = newcap[0:len(s.items)]
}

func (s *Set[K]) minAllocSize(x int) int {
	x -= 1
	x |= x >> 1
	x |= x >> 2
	x |= x >> 4
	x |= x >> 8
	x |= x >> 16
	//最后再加1返回
	return x + 1
}
